package com.citusdata.elven.datastore;

import com.citusdata.elven.datastore.protobuf.Datastore.*;
import com.citusdata.elven.httpRpc.HttpRpcChannel;
import com.citusdata.elven.httpRpc.HttpRpcChannelManager;
import com.citusdata.elven.httpRpc.HttpRpcController;
import com.citusdata.elven.router.Endpoint;
import com.citusdata.elven.router.ReplicationGroupRouter;
import com.google.protobuf.RpcCallback;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;


// Helper class that implements sort and de-duplication order for PutRequests
class PutRequestComparator implements Comparator {
  public int compare(Object first, Object second) {
    PutRequest firstPut  = (PutRequest) first;
    PutRequest secondPut = (PutRequest) second;

    String firstPartition  = firstPut.getPartitionKey();
    String secondPartition = secondPut.getPartitionKey();

    int partitionDiff = firstPartition.compareTo(secondPartition);
    if (partitionDiff != 0) {
      return partitionDiff;
    }

    long firstTimestamp  = firstPut.getTimestamp();
    long secondTimestamp = secondPut.getTimestamp();

    int timestampDiff = (int) (firstTimestamp - secondTimestamp);
    return timestampDiff;
  }
}


public class DatastoreCoordinator {
  private final ReplicationGroupRouter router;
  private final HttpRpcChannelManager  channelManager;
  private final PutRequestComparator   putRequestComparator;

  public DatastoreCoordinator(ReplicationGroupRouter router,
                              HttpRpcChannelManager  channelManager) {
    this.router = router;
    this.channelManager = channelManager;
    this.putRequestComparator = new PutRequestComparator();
  }


  public PutResponse put(PutRequest putRequest) {
    String partitionKey = putRequest.getPartitionKey();
    long timestamp = putRequest.getTimestamp();

    // If any one of the puts fail, we will consider this call to have failed.
    boolean putSucceeded = true;

    List activationTimes = (List) router.getActivationTimes(timestamp, timestamp);
    assert (activationTimes.size() == 1);
    long activationTime = (Long) activationTimes.get(0);
    List endpoints = (List) router.getEndpoints(partitionKey, activationTime);

    for (int ep = 0; ep < endpoints.size(); ++ep) {
      Endpoint endpoint = (Endpoint) endpoints.get(ep);
      String hostname = endpoint.getHostname();
      int port = endpoint.getPort();

      HttpRpcChannel httpChannel = channelManager.getChannel(hostname, port);
      HttpRpcController httpController = new HttpRpcController();

      DatastoreInterface service = DatastoreInterface.newStub(httpChannel);
      DatastoreCallback done = new DatastoreCallback();

      service.put(httpController, putRequest, done);

      channelManager.releaseChannel(httpChannel);
      
      if (httpController.failed()) {
        router.recordFailure(endpoint);
        putSucceeded = false;
      } else {
        PutResponse putResponse = done.putResponse();
        if (!putResponse.getSuccess()) {
          putSucceeded = false;
        }
      }
    }  // for endpoint index

    PutResponse.Builder finalResponse = PutResponse.newBuilder();
    finalResponse.setSuccess(putSucceeded);
    return (finalResponse.build());
  }

  
  public GetRangeResponse getRange(GetRangeRequest getRangeRequest) {
    String partitionKey = getRangeRequest.getPartitionKey();
    long startTime = getRangeRequest.getStartTime();
    long endTime = getRangeRequest.getEndTime();

    // GetRangeResponse contains a list of previously stored put requests. This
    // function retrieves those requests from different endpoints, merges and
    // de-duplicates them.
    Set dedupedPutRequests = new TreeSet(putRequestComparator);
    
    List activationTimes = (List) router.getActivationTimes(startTime, endTime);
    int  lastIndex = activationTimes.size()-1;

    for (int timeIndex = 0; timeIndex <= lastIndex; ++timeIndex) {
      // The router returns a list of activation times from which we retrieve
      // the set of endpoints for that time interval.
      long activationTime = (Long) activationTimes.get(timeIndex);
      int nextIndex = timeIndex+1;
      long nextActivationTime = -1;
      if (nextIndex <= lastIndex) {
        nextActivationTime = (Long) activationTimes.get(nextIndex);
      }
      
      // We decompose an incoming GetRangeRequest into smaller requests, where
      // each request spans over an "activation" interval. We also adjust the
      // start and end times for the first and the last request, respectively.
      long getStartTime = (timeIndex == 0) ? startTime : activationTime;
      long getEndTime   = (nextIndex <= lastIndex) ? nextActivationTime : endTime;

      GetRangeRequest.Builder getBuilder = GetRangeRequest.newBuilder();
      getBuilder.setPartitionKey(partitionKey);
      getBuilder.setStartTime(getStartTime);
      getBuilder.setEndTime(getEndTime);
      GetRangeRequest subGetRangeRequest = getBuilder.build();

      List endpoints = (List) router.getEndpoints(partitionKey, activationTime);
      for (int ep = 0; ep < endpoints.size(); ++ep) {
        Endpoint endpoint = (Endpoint) endpoints.get(ep);
        String hostname = endpoint.getHostname();
        int port = endpoint.getPort();

        HttpRpcChannel httpChannel = channelManager.getChannel(hostname, port);
        HttpRpcController httpController = new HttpRpcController();

        DatastoreInterface service = DatastoreInterface.newStub(httpChannel);
        DatastoreCallback done = new DatastoreCallback();

        service.getRange(httpController, subGetRangeRequest, done);

        channelManager.releaseChannel(httpChannel);

        if (httpController.failed()) {
          router.recordFailure(endpoint);
        } else {
          GetRangeResponse subResponse = done.getRangeResponse();
          List putRequests = subResponse.getPutRequestsList();
          dedupedPutRequests.addAll(putRequests);
        }
      }  // for endpoint index
    }  // for activation time index

    GetRangeResponse.Builder getBuilder = GetRangeResponse.newBuilder();
    getBuilder.addAllPutRequests(dedupedPutRequests);
    return (getBuilder.build());
  }


  private class DatastoreCallback implements RpcCallback {
    PutResponse putResponse;
    GetRangeResponse getRangeResponse;

    public void run(Object object) {
      if (object instanceof PutResponse) {
        putResponse = (PutResponse) object;
      } else if (object instanceof GetRangeResponse) {
        getRangeResponse = (GetRangeResponse) object;
      }
    }

    public PutResponse putResponse() {
      return putResponse;
    }

    public GetRangeResponse getRangeResponse() {
      return getRangeResponse;
    }
  }
}
